/**
 * @author: dn-jinmin/dn-jinmin
 * @doc:
 */

package chatinternal

import (
	"aiworkc/internal/domain"
	"aiworkc/internal/model"
	"aiworkc/internal/svc"
	"aiworkc/pkg/langchain"
	"aiworkc/pkg/langchain/outputparserx"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/prompts"
	"strings"
	"time"
)

type Task struct {
	Type    int
	Title   string
	Content string
}

type ChatLogHandle struct {
	svc    *svc.ServiceContext
	chains chains.Chain
	out    outputparserx.Structured
}

func NewChatLogHandle(svc *svc.ServiceContext) *ChatLogHandle {
	return &ChatLogHandle{
		svc:    svc,
		chains: chains.NewLLMChain(svc.LLMs, prompts.NewPromptTemplate(_defaultChatLogPrompts, []string{"input"})),
		out:    outputparserx.Structured{},
	}
}

func (c *ChatLogHandle) Name() string {
	return "chat_log"
}

func (c *ChatLogHandle) Description() string {
	return "used to summarize and analyze the content of a chat session"

}

func (c *ChatLogHandle) Chains() chains.Chain {
	return chains.NewTransform(c.transform, nil, nil)
}

func (c *ChatLogHandle) transform(ctx context.Context, inputs map[string]any,
	opts ...chains.ChainCallOption) (map[string]any, error) {
	var cid string
	if id, ok := inputs["relationId"].(string); !ok {
		return nil, errors.New("请确定需要总结的会话对象")
	} else {
		cid = id
	}

	startTime, endTime := setTimeRange(inputs)

	msgs, err := c.chatLog(ctx, cid, startTime, endTime)
	if err != nil {
		return nil, err
	}

	// 总结
	res, err := chains.Call(ctx, c.chains, map[string]any{
		"input": msgs,
	})
	if err != nil {
		return nil, err
	}

	text, ok := res[langchain.OutPut].(string)
	if !ok {
		return nil, chains.ErrInvalidOutputValues
	}

	var data []*Task
	if err := json.Unmarshal([]byte(text), &data); err != nil {
		return nil, err
	}

	b, err := json.Marshal(domain.ChatResp{
		ChatType: domain.ChatLog,
		Data:     data,
	})
	if err != nil {
		return nil, err
	}

	return map[string]any{
		langchain.OutPut: string(b),
	}, nil
}

func (c *ChatLogHandle) chatLog(ctx context.Context, cid string, startTime, endTime int64) (string, error) {
	list, err := c.svc.ChatlogModel.ListBySendTime(ctx, cid, startTime, endTime)
	if err != nil {
		return "", err
	}

	// todo
	chatStr := "%s(%s): %s\n"

	var (
		res    strings.Builder
		record = make(map[string]*model.User)
	)

	for i, _ := range list {
		var u *model.User
		if v, ok := record[list[i].SendId]; ok {
			u = v
		} else {
			t, err := c.svc.UserModel.FindOne(ctx, list[i].SendId)
			if err != nil {
				return "", err
			}

			u = t
			record[list[i].SendId] = t
		}

		res.Write([]byte(fmt.Sprintf(chatStr, u.Name, u.ID.Hex(), list[i].MsgContent)))
	}

	return res.String(), nil
}

func setTimeRange(input map[string]any) (int64, int64) {
	var startTime, endTime int64

	cuurentTime := time.Now().Unix()
	if start, ok := input["startTime"].(int); ok && start > 0 {
		startTime = int64(start)
	} else {
		// 获取前一天的消息
		startTime = cuurentTime - 24*3600
	}

	if end, ok := input["endTime"].(int); ok && end > 0 {
		endTime = int64(end)
	} else {
		// 获取前一天的消息
		endTime = cuurentTime
	}

	return startTime, endTime
}
